import datetime
import json
import torch
import re
from typing import List, Dict, Any, Optional, Tuple
from .logger import request_logger

class CitationGenerator:
    """
    引用生成器，用于生成带有引用标记的回答
    基于 deepseek-r1 官方推荐的联网搜索模板
    """
    
    def __init__(self):
        """初始化引用生成器"""
        
        # deepseek-r1 官方推荐的联网搜索模板
        self.search_answer_zh_template = '''# 以下内容是基于用户发送的消息的搜索结果:
{search_results}
在我给你的搜索结果中，每个结果都是[webpage X begin]...[webpage X end]格式的，X代表每篇文章的数字索引。请在适当的情况下在句子末尾引用上下文。请按照引用编号[citation:X]的格式在答案中对应部分引用上下文。如果一句话源自多个上下文，请列出所有相关的引用编号，例如[citation:3][citation:5]，切记不要将引用集中在最后返回引用编号，而是在答案对应部分列出。
在回答时，请注意以下几点：
- 今天是{cur_date}。
- 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。
- 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内，并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项；如非必要，不要主动告诉用户搜索结果未提供的内容。
- 对于创作类的问题（如写论文），请务必在正文的段落中引用对应的参考编号，例如[citation:3][citation:5]，不能只在文章末尾引用。你需要解读并概括用户的题目要求，选择合适的格式，充分利用搜索结果并抽取重要信息，生成符合用户要求、极具思想深度、富有创造力与专业性的答案。你的创作篇幅需要尽可能延长，对于每一个要点的论述要推测用户的意图，给出尽可能多角度的回答要点，且务必信息量大、论述详尽。
- 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在5个点以内，并合并相关的内容。
- 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。
- 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。
- 你的回答应该综合多个相关网页来回答，不能重复引用一个网页。
- 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。

# 用户消息为：
{question}'''
    
    @staticmethod
    def get_used_citations(text: str) -> List[int]:
        """
        从文本中提取所有引用的ID
        
        参数:
            text: 带引用的文本
        
        返回:
            引用ID列表
        """
        # 匹配 [citation:X] 格式的引用
        pattern = r'\[citation:(\d+)\]'
        matches = re.findall(pattern, text)
        # 转换为整数并去重
        citation_ids = list(set([int(m) for m in matches]))
        return sorted(citation_ids)
    
    @staticmethod
    def format_citation_references(answer: str, citations: List[Dict[str, Any]]) -> str:
        """
        在回答末尾添加引用源信息
        
        参数:
            answer: 带引用标记的回答
            citations: 引用源信息列表
            
        返回:
            添加了引用源信息的回答
        """
        if not citations:
            return answer
            
        # 提取使用的引用ID
        used_ids = CitationGenerator.get_used_citations(answer)
        if not used_ids:
            return answer
            
        # 创建引用源信息
        reference_text = "\n\n### 引用来源\n"
        for citation_id in used_ids:
            # 查找对应的引用信息
            citation_info = next((c for c in citations if c.get("id") == citation_id), None)
            if citation_info:
                title = citation_info.get("title", "未知标题")
                source = citation_info.get("source", "知识库")
                url = citation_info.get("url", "")
                
                ref_line = f"[{citation_id}] {title}"
                if source and source != "知识库":
                    ref_line += f" - {source}"
                if url:
                    ref_line += f" [{url}]"
                    
                reference_text += f"{ref_line}\n"
        
        return answer + reference_text
    
    def extract_source_info(self, doc: Any) -> Dict[str, str]:
        """
        从文档中提取源信息
        
        参数:
            doc: 文档对象
            
        返回:
            包含源信息的字典
        """
        source_info = {
            "title": "本地知识库的数据",
            "url": "",
            "source": "知识库",
            "snippet": "",
            "imageUrl": ""
        }
        
        # 提取元数据（如果有）
        if hasattr(doc, 'metadata'):
            metadata = doc.metadata
            
            # 提取标题
            if 'title' in metadata:
                source_info["title"] = metadata['title']
            elif 'name' in metadata:
                source_info["title"] = metadata['name']
            
            # 提取URL
            if 'url' in metadata:
                source_info["url"] = metadata['url']
            elif 'source' in metadata and metadata['source'].startswith('http'):
                source_info["url"] = metadata['source']
            
            # 提取来源
            if 'siteName' in metadata:
                source_info["source"] = metadata['siteName']
            elif 'source' in metadata and not metadata['source'].startswith('http'):
                source_info["source"] = metadata['source']
            
            # 提取图片地址
            if 'imageUrl' in metadata:
                source_info["imageUrl"] = metadata['imageUrl']
        
        # 提取内容片段(最多150个字符)
        if hasattr(doc, 'page_content'):
            content = doc.page_content
            source_info["snippet"] = content[:150] + "..." if len(content) > 150 else content
        
        return source_info
    
    def format_documents_for_citation(self, 
                                      web_docs: List[Any], 
                                      vector_docs: List[Any], 
                                      reranked_docs: List[Any] = None) -> Tuple[str, List[Dict[str, Any]]]:
        """
        将文档格式化为带引用格式的文本
        
        参数:
            web_docs: 网络搜索文档
            vector_docs: 本地向量库文档
            reranked_docs: 重排序后的文档列表（可选）
            
        返回:
            格式化后的文本，适合模型消费
            引用源信息列表
        """
        # 如果提供了重排序后的文档，优先使用
        docs_to_use = reranked_docs if reranked_docs else web_docs + vector_docs
        
        formatted_results = []
        citations_info = []
        
        # 对文档进行编号引用
        for i, doc in enumerate(docs_to_use, 1):
            content = doc.page_content if hasattr(doc, 'page_content') else str(doc)
            metadata = {}
            
            # 提取元数据（如果有）
            if hasattr(doc, 'metadata'):
                metadata = doc.metadata
            
            # 收集引用信息
            source_info = self.extract_source_info(doc)
            source_info["id"] = i  # 添加引用ID
            citations_info.append(source_info)
            
            # 提取来源信息（如果有）
            source = ""
            if 'source' in metadata:
                source = f" (来源: {metadata['source']})"
            elif 'url' in metadata:
                source = f" (来源: {metadata['url']})"
            elif 'siteName' in metadata:
                source = f" (来源: {metadata['siteName']})"
                
            # 格式化为模型期望的格式
            formatted_doc = f"[webpage {i} begin]\n{content}{source}\n[webpage {i} end]"
            formatted_results.append(formatted_doc)
        
        # 连接所有文档
        return "\n\n".join(formatted_results), citations_info
    
    def generate_cited_answer(self, 
                            question: str, 
                            web_docs: List[Any], 
                            vector_docs: List[Any],
                            reranked_docs: List[Any] = None,
                            tokenizer=None,
                            model=None,
                            device=None) -> Tuple[str, List[Dict[str, Any]]]:
        """
        生成带引用的回答
        
        参数:
            question: 用户问题
            web_docs: 网络搜索文档
            vector_docs: 本地向量库文档
            reranked_docs: 重排序后的文档（可选）
            tokenizer: 分词器
            model: 语言模型
            device: 设备（如 'cuda'）
            
        返回:
            带引用的回答文本
            引用源信息列表
        """
        if not model or not tokenizer or not device:
            request_logger.log("生成带引用回答失败：缺少必要的模型组件", level="error")
            return "抱歉，无法生成带引用的回答，因为模型组件未初始化。", []
        
        try:
            # 获取当前日期
            current_date = datetime.datetime.now().strftime("%Y年%m月%d日")
            
            # 格式化搜索结果
            search_results, citations_info = self.format_documents_for_citation(
                web_docs=web_docs,
                vector_docs=vector_docs,
                reranked_docs=reranked_docs
            )
            
            # 如果没有搜索结果，返回提示
            if not search_results:
                return "抱歉，没有找到相关的搜索结果来回答您的问题。", []
            
            # 构建完整的提示
            prompt = self.search_answer_zh_template.format(
                search_results=search_results,
                cur_date=current_date,
                question=question
            )
            
            request_logger.log("生成带引用回答", data={
                "question": question,
                "document_count": len(web_docs) + len(vector_docs),
                "prompt_length": len(prompt)
            })
            
            # 使用模型生成回答
            inputs = tokenizer(prompt, return_tensors="pt").to(device)
            
            # 确保设置正确的注意力掩码，避免警告
            attention_mask = torch.ones_like(inputs["input_ids"])
            inputs["attention_mask"] = attention_mask
            
            # 生成回答
            outputs = model.generate(
                **inputs,
                max_length=4096,  # 更长的最大长度以容纳引用
                min_length=100,
                temperature=0.7,
                top_p=0.9,
                repetition_penalty=1.1,
                no_repeat_ngram_size=3,
                early_stopping=True
            )
            
            # 解码回答
            answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # 提取回答部分，避免包含原始提示
            if "# 用户消息为：" in answer:
                answer = answer.split("# 用户消息为：", 1)[1].strip()
                if question in answer:
                    answer = answer.split(question, 1)[1].strip()
            
            request_logger.log("带引用回答生成完成", data={"answer_length": len(answer)})
            
            # 在回答下方添加引用信息
            answer_with_citations = self.format_citation_references(answer, citations_info)
            
            return answer_with_citations, citations_info
            
        except Exception as e:
            error_msg = f"生成带引用回答出错: {str(e)}"
            request_logger.exception(error_msg)
            return f"抱歉，生成带引用的回答时出现了错误: {str(e)}", []
    
    def generate_stream_cited_answer(self, 
                                   question: str, 
                                   web_docs: List[Any], 
                                   vector_docs: List[Any],
                                   reranked_docs: List[Any] = None,
                                   tokenizer=None,
                                   model=None,
                                   device=None,
                                   stream_controller=None,
                                   user_id=None):
        """
        流式生成带引用的回答
        
        参数:
            question: 用户问题
            web_docs: 网络搜索文档
            vector_docs: 本地向量库文档
            reranked_docs: 重排序后的文档（可选）
            tokenizer: 分词器
            model: 语言模型
            device: 设备（如 'cuda'）
            stream_controller: 流控制器
            user_id: 用户ID
            
        返回:
            生成器，产生回答文本的片段和引用源信息
        """
        if not model or not tokenizer or not device:
            yield f"data: {json.dumps({'error': '模型组件未初始化'}, ensure_ascii=False)}\n\n"
            return
        
        try:
            # 获取当前日期
            current_date = datetime.datetime.now().strftime("%Y年%m月%d日")
            
            # 格式化搜索结果
            search_results, citations_info = self.format_documents_for_citation(
                web_docs=web_docs,
                vector_docs=vector_docs,
                reranked_docs=reranked_docs
            )
            
            # 记录引用信息到日志
            request_logger.log("引用源信息", data={
                "citations_count": len(citations_info),
                "citations": citations_info
            })
            
            # 针对网络搜索的引用源单独记录
            web_citations = [c for c in citations_info if c.get('source') == '网络搜索']
            if web_citations:
                request_logger.log("网络搜索引用源", data={
                    "web_citations_count": len(web_citations),
                    "web_citations": web_citations
                })
            
            # 首先发送引用信息
            yield f"data: {json.dumps({'citations': citations_info}, ensure_ascii=False)}\n\n"
            
            # 如果没有搜索结果，返回提示
            if not search_results:
                yield f"data: {json.dumps({'delta': '抱歉，没有找到相关的搜索结果来回答您的问题。'}, ensure_ascii=False)}\n\n"
                yield "data: [DONE]\n\n"
                return
            
            # 构建完整的提示
            prompt = self.search_answer_zh_template.format(
                search_results=search_results,
                cur_date=current_date,
                question=question
            )
            
            request_logger.log("开始流式生成带引用回答", data={
                "question": question,
                "document_count": len(web_docs) + len(vector_docs),
                "prompt_length": len(prompt)
            })
            
            # 使用模型生成回答
            inputs = tokenizer(prompt, return_tensors="pt").to(device)
            
            # 确保设置正确的注意力掩码，避免警告
            attention_mask = torch.ones_like(inputs["input_ids"])
            inputs["attention_mask"] = attention_mask
            
            # 初始变量
            final_answer = ""
            answer_started = False
            
            # 流式生成
            if hasattr(model, "generate_stream"):
                # 使用原生流式生成（如果支持）
                for output in model.generate_stream(
                    **inputs,
                    max_length=4096,
                    min_length=100,
                    temperature=0.7,
                    top_p=0.9,
                    repetition_penalty=1.1
                ):
                    # 检查流是否被取消
                    if stream_controller and not stream_controller.is_active(user_id):
                        yield "data: [CANCELLED]\n\n"
                        return
                    
                    # 获取当前生成的token
                    if hasattr(output, "token"):
                        token_text = tokenizer.decode([output.token.item()], skip_special_tokens=True)
                    else:
                        token_text = tokenizer.decode(output.sequences[0][-1:], skip_special_tokens=True)
                    
                    # 判断是否开始正式回答
                    if not answer_started:
                        current_text = tokenizer.decode(output.sequences[0], skip_special_tokens=True)
                        if "# 用户消息为：" in current_text and question in current_text:
                            answer_parts = current_text.split(question, 1)
                            if len(answer_parts) > 1 and answer_parts[1].strip():
                                answer_started = True
                                # 开始发送的第一个token
                                first_token = answer_parts[1].strip()
                                if first_token:
                                    final_answer += first_token
                                    if stream_controller:
                                        stream_controller.append_output(user_id, first_token)
                                    yield f"data: {json.dumps({'delta': first_token}, ensure_ascii=False)}\n\n"
                            continue
                    
                    # 如果已经开始回答，累加结果并推送
                    if answer_started and token_text:
                        final_answer += token_text
                        if stream_controller:
                            stream_controller.append_output(user_id, token_text)
                        yield f"data: {json.dumps({'delta': token_text}, ensure_ascii=False)}\n\n"
            else:
                # 备用方案：生成全部文本后拆分为tokens
                outputs = model.generate(
                    **inputs,
                    max_length=4096,
                    min_length=100,
                    temperature=0.7,
                    top_p=0.9,
                    repetition_penalty=1.1,
                    num_return_sequences=1,
                    pad_token_id=tokenizer.eos_token_id,
                    eos_token_id=tokenizer.eos_token_id
                )
                
                full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
                
                # 提取回答部分
                if "# 用户消息为：" in full_text:
                    parts = full_text.split("# 用户消息为：", 1)
                    if len(parts) > 1 and question in parts[1]:
                        answer_text = parts[1].split(question, 1)[1].strip()
                    else:
                        answer_text = parts[1].strip()
                else:
                    answer_text = full_text
                
                # 逐字符输出模拟流式效果
                for char in answer_text:
                    # 检查流是否被取消
                    if stream_controller and not stream_controller.is_active(user_id):
                        yield "data: [CANCELLED]\n\n"
                        return
                    
                    yield f"data: {json.dumps({'delta': char}, ensure_ascii=False)}\n\n"
                    final_answer += char
                    if stream_controller:
                        stream_controller.append_output(user_id, char)
            
            # 流式生成完成，记录结果
            request_logger.log("流式带引用回答生成完成", data={"answer_length": len(final_answer)})
            request_logger.log(f"最终带引用输出内容: {final_answer}")
            
            # 输出完整的引用信息
            request_logger.log("最终所有引用结果", data={
                "all_citations": citations_info,
                "citations_count": len(citations_info)
            })
            
            # 获取使用的引用ID
            used_citation_ids = self.get_used_citations(final_answer)
            request_logger.log("引用使用情况", data={
                "used_citation_ids": used_citation_ids,
                "total_citations": len(citations_info),
                "used_citations_count": len(used_citation_ids)
            })
            
            # 如果有引用，发送引用源信息
            if used_citation_ids:
                # 筛选使用到的引用信息
                used_citations = [c for c in citations_info if c.get("id") in used_citation_ids]
                
                # 记录使用的引用详细信息
                request_logger.log("使用的引用详情", data={
                    "used_citations": used_citations
                })
                
                # 构建引用源信息文本
                reference_text = "\n\n### 引用来源\n"
                for citation_info in used_citations:
                    citation_id = citation_info.get("id")
                    title = citation_info.get("title", "未知标题")
                    source = citation_info.get("source", "知识库")
                    url = citation_info.get("url", "")
                    
                    ref_line = f"[{citation_id}] {title}"
                    if source and source != "知识库":
                        ref_line += f" - {source}"
                    if url:
                        ref_line += f" [{url}]"
                        
                    reference_text += f"{ref_line}\n"
                
                # 发送引用信息
                yield f"data: {json.dumps({'delta': reference_text}, ensure_ascii=False)}\n\n"
                
                # 更新流控制器中的输出
                if stream_controller:
                    stream_controller.append_output(user_id, reference_text)
            
            # 全部生成完成，发送结束标记
            yield "data: [DONE]\n\n"
        
        except Exception as e:
            error_msg = f"生成带引用回答出错: {str(e)}"
            request_logger.exception(error_msg)
            yield f"data: {json.dumps({'error': error_msg}, ensure_ascii=False)}\n\n"
            yield "data: [DONE]\n\n"

# 全局实例
citation_generator = CitationGenerator() 